package com.bigdata.flink.datastream;

import com.bigdata.flink.util.StreamEnvUtil;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoProcessFunction;
import org.apache.flink.util.Collector;

import java.util.*;

public class ConnectStreamTest {
  public static void main(String[] args) throws Exception {
    StreamExecutionEnvironment env = StreamEnvUtil.getEnv();

    DataStreamSource<Tuple2<Integer, String>> source1 = env.fromCollection(Arrays.asList(
        Tuple2.of(1, "java"),
        Tuple2.of(1, "spring"),
        Tuple2.of(2, "hive"),
        Tuple2.of(1, "flink")
    ));

    DataStreamSource<Tuple3<Integer, String, String>> source2 = env.fromCollection(Arrays.asList(
        Tuple3.of(1, "zjx", "teacher"),
        Tuple3.of(2, "gqh", "cto"),
        Tuple3.of(3, "dxj", "leader")
    ));

    // 双流连接，模拟join;数据类型可以不一致，各处理各的;多并行度可以用keyby分区划分数据到统一子任务slot
    ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, String>> connectedStreams = source1.connect(source2);
    ConnectedStreams<Tuple2<Integer, String>, Tuple3<Integer, String, String>> streams = connectedStreams.keyBy(f -> f.f0, f -> f.f0);

    SingleOutputStreamOperator<String> streamOperator = streams.process(new CoProcessFunction<Tuple2<Integer, String>, Tuple3<Integer, String, String>, String>() {
      final Map<Integer, List<Tuple2<Integer, String>>> s1Cache = new HashMap<>();
      final Map<Integer, List<Tuple3<Integer, String, String>>> s2Cache = new HashMap<>();

      @Override
      public void processElement1(Tuple2<Integer, String> value, Context ctx, Collector<String> out) throws Exception {
        Integer key = value.f0;
        if (!s1Cache.containsKey(key)) {
          List<Tuple2<Integer, String>> list = new ArrayList<>();
          list.add(value);
          s1Cache.put(key, list);
        } else {
          s1Cache.get(key).add(value);
        }

        if (s2Cache.containsKey(key)) {
          for (Tuple3<Integer, String, String> s2Element : s2Cache.get(key)) {
            out.collect("s1: " + value + "===== s2: " + s2Element);
          }
        }

      }


      @Override
      public void processElement2(Tuple3<Integer, String, String> value, Context ctx, Collector<String> out) throws Exception {
        Integer key = value.f0;
        if (!s2Cache.containsKey(key)) {
          List<Tuple3<Integer, String, String>> list = new ArrayList<>();
          list.add(value);
          s2Cache.put(key, list);
        } else {
          s2Cache.get(key).add(value);
        }

        if (s1Cache.containsKey(key)) {
          for (Tuple2<Integer, String> s2Element : s1Cache.get(key)) {
            out.collect("s1: " + s2Element + "===== s2: " + value);
          }
        }
      }
    });

    streamOperator.print();
    env.execute();
  }
}
